#结构化提取  Pydantic
import json
from datetime import datetime
from pathlib import Path
from pprint import pprint

from llama_index.core import PromptTemplate
from llama_index.core.output_parsers import PydanticOutputParser
from llama_index.core.program import LLMTextCompletionProgram
from llama_index.readers.file import PDFReader
from pydantic import BaseModel, Field

from llamaIndex.CommonClient import llm
from llamaIndex.class5.Invoice import Invoice

pdf_reader = PDFReader()
documents = pdf_reader.load_data(file=Path("./data/uber_receipt2.pdf"))
text = documents[0].text

sllm = llm.as_structured_llm(Invoice)

response = sllm.complete(text)

json_response = json.loads(response.text)
#解决中文乱码问题
# print(json.dumps(json_response, indent=2, ensure_ascii=False))
#
# pprint(response.raw)

#结构化预测

prompt = PromptTemplate(
    "Extract an invoice from the following text. If you cannot find an invoice ID, use the company name '{company_name}' and the date as the invoice ID: {text}"
)

response = llm.structured_predict(
    Invoice, prompt, text=text, company_name="Uber"
)
json_output = response.model_dump_json()
# print(json.dumps(json.loads(json_output), indent=2, ensure_ascii=False))


textCompletion = LLMTextCompletionProgram.from_defaults(
    output_cls=Invoice,
    llm=llm,
    prompt=PromptTemplate(
        "Extract an invoice from the following text. If you cannot find an invoice ID, use the company name '{company_name}' and the date as the invoice ID: {text}"
    ),
)

output = textCompletion(company_name="Uber", text=text)
json_output2 = output.model_dump_json()
# print(json.dumps(json.loads(json_output2), indent=2, ensure_ascii=False))

class MyOutputParser(PydanticOutputParser):
    def get_pydantic_object(self, text: str):
        # do something more clever than this
        return self.output_parser.model_validate_json(text)


textCompletion = LLMTextCompletionProgram.from_defaults(
    llm=llm,
    prompt=PromptTemplate(
        "Extract an invoice from the following text. If you cannot find an invoice ID, use the company name '{company_name}' and the date as the invoice ID: {text}"
    ),
    output_parser=MyOutputParser(output_cls=Invoice),
)
output = textCompletion(company_name="Uber", text=text)
json_output3 = output.model_dump_json()
print(json.dumps(json.loads(json_output3), indent=2, ensure_ascii=False))